#include <pthread.h>
#include <mutex>
#include <thread>
#include <unistd.h>
#include <signal.h>
#include <mysql.h>
#include <map>
#include "opencv2/opencv.hpp"

#include "cnn_function.hpp"
#include "util.h"

#include "rpc/server.h"

std::vector<std::string> current_user_name;
std::vector<int> current_user_spoofing;
std::vector<cv::Rect> rois_global;

// 0.62 for 50, 0.6 for 100
// sudo apt-get install libmysqlclient-dev
std::map<int, float**> all_group_feature;
std::map<int, int*> all_group_feature_id;
std::map<int, int> all_group_index;
std::map<int, float> group_id_to_threshold;

bool have_init = false;
std::mutex gpu_lock;
bool recognition_stop = false;

int reload_feature();
int load_group_id_to_threshold();

void free_dnn_network(){
    uninit_detector();
    uninit_recognition();
    uninit_mtcnn();
}

void init_dnn_network(){
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
    have_init = true;
    init_detector("/home/iim/.bin/git/cnn/cfg/yolov3-small.cfg", "/home/iim/.bin/git/cnn/model/yolov3-small_final_max_epoch_15.weights");
    init_recognition("/home/iim/.bin/git/cnn/cfg/cosface_new.cfg", "/home/iim/.bin/git/cnn/model/model.cnn.50");
    init_mtcnn("/home/iim/.bin/git/cnn/cfg/mtcnn_onet.cfg", "/home/iim/.bin/git/cnn/model/mtcnn_final.weights");
    reload_feature();
    //load_group_id_to_threshold();
}

int set_threshold(int group_id, float threshold) {
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
    group_id_to_threshold[group_id] = threshold;
    printf("set_threshold %f\n", threshold);
    return 1;
}

void print_error_message(MYSQL *connection){
    if (mysql_errno(connection)) {
        fprintf(stderr, "mysql error %d: %s\n", mysql_errno(connection), mysql_error(connection));
    }
}

int load_group_id_to_threshold(){
    group_id_to_threshold.clear();
    MYSQL *connection = mysql_init(NULL);
    if (connection == NULL) {
        printf("mysql_init failed!\n");
        return -1;
    }
    connection = mysql_real_connect(connection, "localhost", "iim", "mysql@120test", "recognition", 0, NULL, 0);
    if (connection) {
        printf("Connection success! ");
    } else {
        print_error_message(connection);
        mysql_close(connection);
        return 1;
    }
    char query_sql[256] = {0};
    sprintf(query_sql, "select * from `group` where deleted = 0 order by id");
    int query_result = mysql_query(connection, query_sql);
    if(query_result){
        print_error_message(connection);
        mysql_close(connection);
        return 2;
    } else {
        MYSQL_RES *res_ptr = mysql_store_result(connection);
        if(res_ptr) {
            printf("query %lu Rows", (unsigned long)mysql_num_rows(res_ptr));
            unsigned int num_fields = mysql_num_fields(res_ptr);
            printf("num_fields: %u\n", num_fields);
            MYSQL_ROW row;
            while((row = mysql_fetch_row(res_ptr)))  {
                int group_id = -1;
                float threshold = -1;
                for(unsigned int i = 0; i < num_fields; i++){
                    //printf("%s ",row[i]);
                    if(i == 0){
                        group_id = atoi(row[i]);
                    } else if (3 == i){
                        threshold = atof(row[i]);
                    } else {
                    }
                }
                group_id_to_threshold[group_id] = threshold;
            }
            if (mysql_errno(connection)) {
                fprintf(stderr, "Retrive error:%s\n", mysql_error(connection));
            }
            mysql_free_result(res_ptr);
        } else {
            print_error_message(connection);
            mysql_close(connection);
            return 3;
        }
    }
    mysql_close(connection);
    std::map<int, float>::iterator it;
    for(it = group_id_to_threshold.begin(); it != group_id_to_threshold.end(); it++){
        printf("group_id_to_threshold: %d %f\n", it->first, it->second);
    }
    return 0;
}

int reload_feature() {
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
    all_group_feature.clear();
    all_group_feature_id.clear();
    MYSQL *connection = mysql_init(NULL);
    if (connection == NULL) {
        printf("mysql_init failed!\n");
        return -1;
    }
    connection = mysql_real_connect(connection, "localhost", "iim", "mysql@120test", "recognition", 0, NULL, 0);
    if (connection) {
        printf("Connection success! ");
    } else {
        print_error_message(connection);
        mysql_close(connection);
        return 1;
    }
    char query_sql[256] = {0};
    sprintf(query_sql, "select * from feature where status = 1 order by id");
    int query_result = mysql_query(connection, query_sql);
    if(query_result){
        print_error_message(connection);
        mysql_close(connection);
        return 2;
    } else {
        MYSQL_RES *res_ptr = mysql_store_result(connection);
        if(res_ptr) {
            unsigned long row_num = mysql_num_rows(res_ptr)
            printf("query %lu Rows", row_num);
            unsigned int num_fields = mysql_num_fields(res_ptr);
            printf("num_fields: %u\n", num_fields);
            MYSQL_ROW row;
            while((row = mysql_fetch_row(res_ptr)))  {
                //unsigned long *field_lengths = mysql_fetch_lengths(res_ptr);
                int feature_id = -1;
                int group_id = -1;
                float *feature_float = malloc(512 * sizeof(float));
                for(unsigned int i = 0; i < num_fields; i++){
                    //printf("%s ",row[i]);
                    if(i == 0){
                        feature_id = atoi(row[i]);
                    } else if (5 == i){
                        group_id = atoi(row[i]);
                    }else if(8 == i){
                        char *float_str = (char *)(row[i]);
                        for(unsigned int j = 0; j < 512; j++) {
                            feature_float[j] = atof(float_str);
                            float_str = strchr(float_str, ',') + 1;
                        }
                    } else {
                    }
                }
                std::map<int, std::vector<int>>::iterator it = all_group_feature_id.find(group_id);
                if (it == all_group_feature_id.end()){
                    all_group_feature_float_ptr = calloc(row_num, sizeof(float **));
                    all_group_feature[group_id] = all_group_feature_float_ptr;
                    all_group_feature_id_ptr = calloc(row_num, sizeof(int *));
                    all_group_feature_id[group_id] = all_group_feature_id_ptr;
                    all_group_index[group_id] = 0;
                }
                all_group_feature_id[group_id].push_back(feature_id);
                all_group_feature[group_id].push_back(feature_float);
            }
            if (mysql_errno(connection)) {
                fprintf(stderr, "Retrive error:%s\n", mysql_error(connection));
            }
            mysql_free_result(res_ptr);
        } else {
            print_error_message(connection);
            mysql_close(connection);
            return 3;
        }
    }
    mysql_close(connection);

    
    std::map<int, std::vector<int>>::iterator it;
    for(it = all_group_feature_id.begin(); it != all_group_feature_id.end(); it++){
        for(unsigned int i = 0; i < all_group_feature_id[it->first].size(); i++){
            printf("%d: %d %d\n", i, it->first, it->second[i]);
        }
    }

    std::map<int, std::vector<std::vector<float>>>::iterator itt;
    for(itt = all_group_feature.begin(); itt != all_group_feature.end(); itt++){
        for(unsigned int i = 0; i < all_group_feature[itt->first].size(); i++){
            printf("%d: %d %f %f %f %f\n", i, itt->first, itt->second[i][0], itt->second[i][1], itt->second[i][510], itt->second[i][511]);
        }
    }
    
    printf("all_group_feature size %lu\n", all_group_feature.size());
    printf("reload_feature over\n");
    return 0;
}

std::string registration_image(std::vector<char> const& img) {
    double start = getMsecOfNow();
    cv::Mat image_input = cv::imdecode(img, CV_LOAD_IMAGE_COLOR);
    std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);

    int code = 0;
    if(!have_init){
        init_dnn_network();
    }
    int detection_bbox[MAX_BBOX_NUM * 4];

    std::vector<cv::Mat> aligned_faces;
    int face_features_length = 512;
    int face_count = 0;
    detection_face(image_input, &face_count, detection_bbox);
    if(face_count == 0){
        code = 1;  // no face
        return std::to_string(code);
    } else if(face_count > 1){
        code = 2;  // multi face
        return std::to_string(code);
    }
    if(detection_bbox[2] - detection_bbox[0] < 10 && detection_bbox[3] - detection_bbox[1] < 10){
        code = 3;  // face size small
        return std::to_string(code);
    }
    float *face_feature = (float *)malloc(face_count * face_features_length *sizeof(float));
    extract_face_feature(image_input, face_count, face_feature, detection_bbox, &code, 1);
    std::string registration_image_feature_str;
    if(code == 0){
        std::string feature_str = "";
        for(int i = 0; i < 512; i++){
            feature_str += std::to_string(face_feature[i]) + ",";
        }
        registration_image_feature_str = feature_str;
    } else {
        registration_image_feature_str = std::to_string(code);
    }
    free(face_feature);
    printf("registration_image, face_count %d image width %d, height %d, over spend: %f ms\n",
           face_count, image_input.cols, image_input.rows, getMsecOfNow() - start);
    std::cout << "registration_image spend: " << getMsecOfNow() - start << " ms code " << code << ", face_count " << face_count
              << " image width " << image_input.cols << " image width " << image_input.rows
              <<  " thread_id " <<  std::this_thread::get_id() << std::endl;
    return registration_image_feature_str;
}

std::string recognition_face(std::vector<char> const& img, int group_id) {
    if(!have_init){
        init_dnn_network();
    }
    int code = 0;
    int face_count = 0;
    std::string bbox_str = "";
    std::string feature_ids = "";
    std::map<int, std::vector<int>>::iterator all_group_feature_id_it = all_group_feature_id.find(group_id);
    printf("well\n");
    if (all_group_feature_id_it == all_group_feature_id.end()){
        code = 2;  // not registration
        return std::to_string(code) + "|" + std::to_string(face_count) + "|" + feature_ids + "|" + bbox_str;
    }
    int detection_bbox[MAX_BBOX_NUM * 4];
    cv::Mat image_input = cv::imdecode(img, CV_LOAD_IMAGE_COLOR);
    double start = getMsecOfNow();
    int face_features_length = 512;
    float *face_feature = (float *)malloc(face_count * face_features_length * sizeof(float));
    {
        std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);

        printf("well 0\n");
        detection_face(image_input, &face_count, detection_bbox);
        for(int i = 0; i < 4; i++) printf("%d %d\n", i, detection_bbox[i]);
        printf("well 0\n");
        if(face_count == 0){
            code = 1;  // no face
            free(face_feature);
            return std::to_string(code) + "|" + std::to_string(face_count) + "|" + feature_ids + "|" + bbox_str;
        }
        printf("well 3 %d\n", face_count);
        extract_face_feature(image_input, face_count, face_feature, detection_bbox, &code, 0);
    }
    printf("well 1\n");
    std::vector<std::vector<float>> &group_feature = all_group_feature[group_id];
    unsigned int group_feature_num = group_feature.size();
    for(int j = 0; j < face_count; j++){
        float max_score = -1.0F;
        int max_score_feature_id = -1;
        for(unsigned int i = 0; i < group_feature_num; i++){
            float score = 0;
            std::vector<float> &register_feature = group_feature[i];
            for(int m = 0; m < face_features_length; m++){
                score += face_feature[j * face_features_length + m] * register_feature[m];
            }
            if(score > max_score){
                max_score = score;
                max_score_feature_id = i;
            }
        }
        printf("max_score %f\n", max_score);
        if(max_score > 0.6){ //group_id_to_threshold[group_id]){
            feature_ids += std::to_string(all_group_feature_id[group_id][max_score_feature_id]) + ",";
            for(int i = 0; i < 4; i++){
                bbox_str += std::to_string(detection_bbox[j * 4 + i]) + ",";
            }
        }

    }
    std::cout << "recognition_face spend: " << getMsecOfNow() - start << " ms code, face_count " << face_count
              <<  " thread_id" <<  std::this_thread::get_id() << std::endl;
    free(face_feature);
    return std::to_string(code) + "|" + std::to_string(face_count) + "|" + feature_ids + "|" + bbox_str;
}

void sig_close(int sig){
    const char *s;
    switch (sig) {
    case SIGINT:
        s = "interrupt";
        break;
    case SIGTERM:
        s = "terminate";
        break;
    }
    recognition_stop = true;
    printf("\n\nreceived signal: %s, quit...\n", s);
}

int main(int argc, char **argv)
{
    signal(SIGINT, sig_close);  //close CTRL+C
    signal(SIGTERM, sig_close);

    rpc::server srv("127.0.0.1", 18080);
    srv.bind("registration_image", &registration_image);
    srv.bind("recognition_face", &recognition_face);
    srv.bind("set_threshold", set_threshold);
    srv.bind("reload_feature", reload_feature);
    //srv.run();
    srv.async_run();
    while (!recognition_stop){
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    printf("rpc quit...\n");
    if(have_init){
        free_dnn_network();
    }
    return 0;
}
